"""
把 *任何* nn.Module → FedCP 期望的接口
若模型本身已经有 .head / .feature_extractor 就什么也不做
"""
from torch import nn


class _Identity(nn.Module):          # 只为占位
    def forward(self, x): return x


def wrap(model: nn.Module) -> nn.Module:
    # 如果作者在模型里已经做了①②步，就直接返回
    if hasattr(model, "head") and hasattr(model, "feature_extractor"):
        return model

    # ----- 自动分离最后一个 Linear 当分类头 -----
    # （适用于大多数 PyTorch 模型；真的找不到时，需要手动在模型里加上 ①②）
    head = None
    for name, module in reversed(list(model.named_children())):
        if isinstance(module, nn.Linear):
            head = module
            setattr(model, name, _Identity())   # backbone 失去分类功能
            break
    if head is None:
        raise ValueError("⚠️ 未找到可拆的 Linear 作为 head，请手动在模型里设置 self.head")

    class _FedCP(nn.Module):
        def __init__(self, backbone, head_layer):
            super().__init__()
            self.backbone = backbone  # 整体骨干
            self.head = head_layer  # 分类头

            # === ★ 关键改动：保证 feature_extractor 是 nn.Module ===
            # ① 如果骨干自己带 feature_extractor，就用它
            # ② 否则直接把骨干当特征提取器（而不是 backbone.forward）
            self.feature_extractor = (
                backbone.feature_extractor
                if hasattr(backbone, "feature_extractor")
                else backbone  # ← nn.Module，而非函数
            )

        def forward(self, *args, **kwargs):
            feat = self.feature_extractor(*args, **kwargs)  # 依然可调用
            return self.head(feat)